from mmdet.registry import DATASETS
import pickle
from mmengine.config import Config
import numpy as np
from mmdet.evaluation import eval_map
import matplotlib.pyplot as plt
import matplotlib

# 设置matplotlib使用中文字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题

# 加载测试结果和数据集
with open('/media/ross/8TB/project/lsh/deep_learning/DiffusionDet_mmdet/DiffusionDet/work_dirs/ablation/diffusiondet_r50_lamfpn8_epoch_microalgeaOri_0lcm2_1adem2_1ddim4_1distill4_config2/20250423_210113/result.pkl', 'rb') as f:
    results = pickle.load(f)
cfg = Config.fromfile(
    '/media/ross/8TB/project/lsh/deep_learning/DiffusionDet_mmdet/DiffusionDet/work_dirs/ablation/diffusiondet_r50_lamfpn8_epoch_microalgeaOri_0lcm2_1adem2_1ddim4_1distill4_config2/diffusiondet_r50_lamfpn8_epoch_microalgeaOri_0lcm2_1adem2_1ddim4_1distill4_config2.py')

# 检查结果格式
print("\n检查结果格式:")
print(f"结果类型: {type(results)}")
print(f"结果长度: {len(results)}")
if len(results) > 0:
    print(f"第一个结果类型: {type(results[0])}")
    if isinstance(results[0], dict):
        print(f"第一个结果的键: {list(results[0].keys())}")
    elif isinstance(results[0], list) or isinstance(results[0], tuple):
        print(f"第一个结果的长度: {len(results[0])}")
        if len(results[0]) > 0:
            print(f"第一个结果中第一个元素的类型: {type(results[0][0])}")
            if isinstance(results[0][0], np.ndarray):
                print(f"第一个结果中第一个元素的形状: {results[0][0].shape}")

# 使用 DATASETS.build() 替代旧版的 build_dataset
dataset = DATASETS.build(cfg.test_dataloader.dataset)

# 获取每个类别的AP
class_names = dataset.metainfo['classes']
print(f"\n类别名称: {class_names}")
print(f"类别数量: {len(class_names)}")

# 提取标注信息
annotations = []
for i in range(len(dataset)):
    data_info = dataset.get_data_info(i)
    instances = data_info.get('instances', [])

    bboxes = []
    labels = []

    for instance in instances:
        if instance.get('ignore_flag', 0) == 0:  # 不忽略的实例
            bboxes.append(instance['bbox'])
            labels.append(instance['bbox_label'])

    # 将标注信息转换为 numpy 数组
    ann = {
        'bboxes': np.array(bboxes, dtype=np.float32),
        'labels': np.array(labels, dtype=np.int64)
    }
    annotations.append(ann)

# 检查是否成功提取标注
print(f"\n提取了 {len(annotations)} 个样本的标注")
print(f"第一个样本有 {len(annotations[0]['bboxes'])} 个边界框")
print(f"边界框示例: {annotations[0]['bboxes'][:2]}")
print(f"标签示例: {annotations[0]['labels'][:2]}")

# 根据结果格式调整评估方法
num_classes = len(class_names)

# 如果结果是字典格式，需要转换为列表格式
if isinstance(results[0], dict) and 'pred_instances' in results[0]:
    print("\n检测到结果是字典格式，正在转换...")
    converted_results = []
    for result in results:
        pred_instances = result['pred_instances']
        scores = pred_instances.get('scores', []).cpu().numpy()
        bboxes = pred_instances.get('bboxes', []).cpu().numpy()
        labels = pred_instances.get('labels', []).cpu().numpy()

        # 按类别整理结果
        result_list = [np.zeros((0, 5), dtype=np.float32) for _ in range(num_classes)]
        for i, (bbox, score, label) in enumerate(zip(bboxes, scores, labels)):
            if label < num_classes:  # 确保标签在有效范围内
                # 将 bbox 和 score 合并为 [x1, y1, x2, y2, score] 格式
                det = np.concatenate([bbox, np.array([score])], axis=0)
                result_list[label] = np.vstack([result_list[label], det])

        converted_results.append(result_list)

    results = converted_results
    print(f"转换后的结果长度: {len(results)}")
    if len(results) > 0:
        print(f"第一个结果的长度: {len(results[0])}")
        for i, cls_result in enumerate(results[0]):
            print(f"类别 {i} 的检测结果数量: {len(cls_result)}")

# 用于存储每个类别的AP值
class_ap_values = []
class_names_en = []  # 英文类别名称
class_recall_values = []  # 用于存储每个类别的Recall值

# 尝试计算每个类别的 AP
try:
    print("\n开始评估每个类别的 AP...")
    mean_ap, ap_results = eval_map(
        results,
        annotations,
        iou_thr=0.5
    )

    # 打印结果
    print("\n每个类别的 AP 结果:")
    for idx, cls_ap in enumerate(ap_results):
        ap_value = cls_ap['ap']
        recall_value = cls_ap['recall'][-1] if len(cls_ap['recall']) > 0 else 0
        class_ap_values.append(ap_value)
        class_names_en.append(class_names[idx])  # 使用原始英文类别名称
        class_recall_values.append(recall_value)
        print(f"Class {class_names[idx]} (ID:{idx}) | AP@0.5 = {ap_value:.4f}")
        print(f"   Recall@0.5 = {recall_value:.2f}")

    # 计算 mAP
    print(f"\nmAP@0.5 = {mean_ap:.4f}")

except Exception as e:
    print(f"\n评估时出错: {e}")
    print("\n尝试手动计算每个类别的 AP...")

    # 手动计算每个类别的 AP
    class_aps = []
    for class_id in range(num_classes):
        print(f"\n计算类别 {class_names[class_id]} (ID:{class_id}) 的 AP...")

        # 为当前类别创建单独的结果和标注
        class_results = []
        class_annotations = []

        for i, ann in enumerate(annotations):
            # 过滤当前类别的标注
            mask = ann['labels'] == class_id
            class_ann = {
                'bboxes': ann['bboxes'][mask] if len(mask) > 0 else np.zeros((0, 4), dtype=np.float32),
                'labels': np.zeros(np.sum(mask), dtype=np.int64)  # 单类别评估时，标签都设为 0
            }
            class_annotations.append(class_ann)

            # 提取当前类别的预测结果
            if isinstance(results[i], list) and len(results[i]) > class_id:
                class_results.append([results[i][class_id]])  # 注意：需要包装在列表中
            else:
                class_results.append([np.zeros((0, 5), dtype=np.float32)])  # 空结果

        # 计算当前类别的 AP
        try:
            class_mean_ap, class_ap_result = eval_map(
                class_results,
                class_annotations,
                iou_thr=0.5
            )
            ap_value = class_mean_ap
            recall_value = class_ap_result[0]['recall'][-1] if class_ap_result and len(
                class_ap_result[0]['recall']) > 0 else 0
            class_ap_values.append(ap_value)
            class_names_en.append(class_names[class_id])  # 使用原始英文类别名称
            class_recall_values.append(recall_value)

            class_aps.append({
                'class_id': class_id,
                'class_name': class_names[class_id],
                'ap': ap_value,
                'recall': recall_value,
                'detail': class_ap_result[0] if class_ap_result else None
            })
            print(f"类别 {class_names[class_id]} (ID:{class_id}) | AP@0.5 = {ap_value:.4f}")
        except Exception as e:
            print(f"计算类别 {class_names[class_id]} 的 AP 时出错: {e}")
            class_ap_values.append(0.0)
            class_names_en.append(class_names[class_id])  # 使用原始英文类别名称
            class_recall_values.append(0.0)

            class_aps.append({
                'class_id': class_id,
                'class_name': class_names[class_id],
                'ap': 0.0,
                'recall': 0.0,
                'detail': None
            })

    # 打印每个类别的 AP
    print("\n每个类别的 AP 结果:")
    for cls_ap in class_aps:
        print(f"Class {cls_ap['class_name']} (ID:{cls_ap['class_id']}) | AP@0.5 = {cls_ap['ap']:.4f}")
        if cls_ap['detail'] is not None:
            recall = cls_ap['recall']
            print(f"   Recall@0.5 = {recall:.2f}")

    # 计算 mAP
    mean_ap = sum([cls_ap['ap'] for cls_ap in class_aps]) / len(class_aps)
    print(f"\nmAP@0.5 = {mean_ap:.4f}")

# 创建图表1：每个类别的AP值
plt.figure(figsize=(14, 6))
bars = plt.bar(class_names_en, class_ap_values, color='skyblue')
plt.axhline(y=mean_ap, color='r', linestyle='-', label=f'mAP = {mean_ap:.4f}')
plt.xlabel('Class', fontsize=12)
plt.ylabel('AP@0.5', fontsize=12)
plt.title('AP@0.5 for Each Class', fontsize=14)
plt.xticks(rotation=45, ha='right')
plt.ylim(0, 1.0)
plt.grid(axis='y', linestyle='--', alpha=0.7)

# 在每个柱子上方添加AP值
for bar, ap_value in zip(bars, class_ap_values):
    plt.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
             f'{ap_value:.3f}', ha='center', fontsize=9)

plt.tight_layout()
plt.legend()
plt.savefig(
    '/media/ross/8TB/project/lsh/deep_learning/DiffusionDet_mmdet/DiffusionDet/tools/mymodel_analysis_tools/class_ap_chart.png',
    dpi=300)
plt.show()

# 创建图表2：每个类别的AP和Recall值对比
plt.figure(figsize=(14, 6))
x = np.arange(len(class_names_en))
width = 0.35

bars1 = plt.bar(x - width / 2, class_ap_values, width, label='AP@0.5', color='skyblue')
bars2 = plt.bar(x + width / 2, class_recall_values, width, label='Recall@0.5', color='lightgreen')

plt.xlabel('Class', fontsize=12)
plt.ylabel('Score', fontsize=12)
plt.title('AP@0.5 and Recall@0.5 for Each Class', fontsize=14)
plt.xticks(x, class_names_en, rotation=45, ha='right')
plt.ylim(0, 1.0)
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.legend()

# 在每个柱子上方添加值
for bar, value in zip(bars1, class_ap_values):
    plt.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
             f'{value:.2f}', ha='center', fontsize=8)
for bar, value in zip(bars2, class_recall_values):
    plt.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
             f'{value:.2f}', ha='center', fontsize=8)

plt.tight_layout()
plt.savefig(
    '/media/ross/8TB/project/lsh/deep_learning/DiffusionDet_mmdet/DiffusionDet/tools/mymodel_analysis_tools/class_ap_recall_chart.png',
    dpi=300)
plt.show()

# 创建图表3：AP和Recall的散点图
plt.figure(figsize=(10, 8))
plt.scatter(class_ap_values, class_recall_values, color='blue', s=100, alpha=0.7)

# 添加类别标签
for i, txt in enumerate(class_names_en):
    plt.annotate(txt, (class_ap_values[i], class_recall_values[i]),
                 xytext=(5, 5), textcoords='offset points', fontsize=10)

plt.xlabel('AP@0.5', fontsize=12)
plt.ylabel('Recall@0.5', fontsize=12)
plt.title('AP@0.5 vs Recall@0.5 for Each Class', fontsize=14)
plt.grid(True, linestyle='--', alpha=0.7)
plt.xlim(0, 1.0)
plt.ylim(0, 1.0)

# 添加对角线参考
plt.plot([0, 1], [0, 1], 'r--', alpha=0.5)

plt.tight_layout()
plt.savefig(
    '/media/ross/8TB/project/lsh/deep_learning/DiffusionDet_mmdet/DiffusionDet/tools/mymodel_analysis_tools/ap_recall_scatter.png',
    dpi=300)
plt.show()
